from options.base_options import BaseOptions


class TrainOptions(BaseOptions):
    def __init__(self):
        super(TrainOptions, self).__init__()
        parser = self.init(self.parser)
        opt, parser = self.gather_options(parser)
        self.opt = opt
        self.parser = parser

    def init(self, parser):
        # freq
        parser.add_argument('--display_freq', type=int, default=1000, help='frequency of showing training results on screen')
        parser.add_argument('--print_freq', type=int, default=100, help='frequency of showing training results on console')
        parser.add_argument('--save_latest_freq', type=int, default=5000, help='frequency of saving the latest results')
        parser.add_argument('--save_epoch_freq', type=int, default=10, help='frequency of saving checkpoints at the end of epochs')

        # training
        parser.add_argument('--continue_train', action='store_true', help='continue training: load the latest model')
        parser.add_argument('--which_epoch', type=str, default='latest',
                            help='which epoch to load? set to latest to use latest cached model')
        parser.add_argument('--niter', type=int, default=100, help='# of iter at starting learning rate. This is NOT the total #epochs. Totla #epochs is niter + niter_decay')
        parser.add_argument('--niter_decay', type=int, default=100, help='# of iter to linearly decay learning rate to zero')
        parser.add_argument('--optimizer', type=str, default='adam')
        parser.add_argument('--beta1', type=float, default=0.5, help='momentum term of adam')
        parser.add_argument('--beta2', type=float, default=0.999, help='momentum term of adam')
        parser.add_argument('--lr', type=float, default=0.0002, help='initial learning rate for adam')
        parser.add_argument('--D_steps_per_G', type=int, default=1, help='number of discriminator iterations per generator iterations.')
        parser.add_argument('--use_ema', action='store_true', help='if true, use EMA in G')
        parser.add_argument('--ema_beta', type=float, default=0.999, help='beta in ema setting')

        # norm type
        parser.add_argument('--norm_G', type=str, default='spectralinstance', help='instance normalization(IN) or batch normalization(BN)')
        parser.add_argument('--norm_D', type=str, default='spectralinstance', help='instance normalization(IN) or batch normalization(BN)')
        parser.add_argument('--norm_E', type=str, default='spectralinstance', help='instance normalization(IN) or batch normalization(BN)')

        # spade params
        parser.add_argument('--label_nc', type=int, default=1,
                            help='segmantic input channel num')
        # parser.add_argument('--contain_dontcare_label', action='store_true',
        #                     help='if the label map contains dontcare label (dontcare=255)')

        parser.add_argument('--output_nc', type=int, default=3, help='# of output image channels')
        parser.add_argument('--eqlr_sn', action='store_true', help='if true, use equlr, else use sn')
        parser.add_argument('--PONO', action='store_true', help='use positional normalization ')
        parser.add_argument('--PONO_C', action='store_true', help='use C normalization in corr module')
        parser.add_argument('--use_se', action='store_true', help='use SE layer in spade')
        parser.add_argument('--norm_type', type=str, default='BN', help='spade normalization type, choose in BN(batch norm), IN(instance norm)', choices=('IN', 'BN'))

        # for generator and discriminator
        parser.add_argument('--ngf', type=int, default=64, help='# of discrim filters in first conv layer')
        parser.add_argument('--use_attention', action='store_true', help='and nonlocal block in G and D')
        parser.add_argument('--no_ganFeat_loss', action='store_true', help='if specified, do *not* use discriminator feature matching loss')
        parser.add_argument('--gan_mode', type=str, default='hinge', help='(ls|original|hinge)')
        parser.add_argument('--netD', type=str, default='multiscale', help='(n_layers|multiscale|image)')
        parser.add_argument('--no_TTUR', action='store_true', help='Use TTUR training scheme')
        parser.add_argument('--n_layers_D', type=int, default=4, help='# layers in each discriminator')
        parser.add_argument('--num_D', type=int, default=2, help='number of discriminators to be used in multiscale')
        parser.add_argument('--use_22ctx', action='store_true', help='if true, also use 2-2 in ctx loss')
        parser.add_argument('--ctx_w', type=float, default=1.0, help='ctx loss weight')

        parser.add_argument('--netG', type=str, default='spade', help='selects model to use for netG (pix2pixhd | spade)')
        parser.add_argument('--init_type', type=str, default='xavier', help='networks initialization [normal|xavier|kaiming|orthogonal]')
        parser.add_argument('--init_variance', type=float, default=0.02, help='variance of the initialization distribution')
        parser.add_argument('--z_dim', type=int, default=256, help="dimension of the latent z vector")
        parser.add_argument('--D_type', type=str, default='structure', help='decide structure discriminator or color discriminator')

        # for correspondence
        parser.add_argument('--warp_patch', action='store_true', help='use corr matrix to warp 4*4 patch')
        parser.add_argument('--warp_stride', type=int, default=4, help='corr matrix 256 / warp_stride')
        parser.add_argument('--weight_domainC', type=float, default=0.0, help='weight of Domain classification loss for domain adaptation')
        parser.add_argument('--domain_rela', action='store_true', help='if true, use Relativistic loss in domain classifier')
        parser.add_argument('--adaptor_kernel', type=int, default=3, help='kernel size in domain adaptor')
        # parser.add_argument('--adaptor_se', action='store_true', help='if true, use se layer in domain adaptor')
        parser.add_argument('--adaptor_res_deeper', action='store_true', help='if true, use 6 res block in domain adaptor')
        parser.add_argument('--adaptor_nonlocal', action='store_true', help='if true, use nonlocal block in domain adaptor')
        parser.add_argument('--warp_cycle_w', type=float, default=0.0, help='push warp cycle to ref')
        parser.add_argument('--warp_self_w', type=float, default=0.0, help='push warp self to ref')
        parser.add_argument('--two_cycle', action='store_true', help='input to ref and back')
        parser.add_argument('--match_kernel', type=int, default=3, help='correspondence matrix match kernel size')
        parser.add_argument('--show_warpmask', action='store_true', help='save warp mask')
        parser.add_argument('--CBN_intype', type=str, default='warp_mask', help='type of CBN input for framework, warp/mask/warp_mask')

        # for vgg feature net
        parser.add_argument('--which_perceptual', type=str, default='5_2', help='relu5_2 or relu4_2')
        parser.add_argument('--weight_perceptual', type=float, default=0.01)
        parser.add_argument('--vgg_normal_correct', action='store_true', help='if true, correct vgg normalization and replace vgg FM model with ctx model')
        parser.add_argument('--novgg_featpair', type=float, default=10.0, help='in no vgg setting, use pair feat loss in domain adaptation')
        parser.add_argument('--dilation_conv', action='store_true',
                            help='if true, use dilation conv in domain adaptor when adaptor_res_deeper is True')
        parser.add_argument('--use_coordconv', action='store_true', help='if true, use coordconv in CorrNet')
        parser.add_argument('--fm_ratio', type=float, default=0.1, help='vgg fm loss weight comp with ctx loss')

        # for dataset mode
        parser.add_argument('--hard_reference_probability', type=float, default=0.5, help='hard reference training probability')

        # for loss
        parser.add_argument('--weight_gan', type=float, default=10.0, help='weight of all gan loss')
        parser.add_argument('--weight_ctx', type=float, default=1.0, help='ctx loss weight')
        parser.add_argument('--D_cam', type=float, default=0.0, help='weight of CAM loss in D')
        parser.add_argument('--warp_mask_losstype', type=str, default='none', help='type of warped mask loss, none/direct/cycle')
        parser.add_argument('--lambda_feat', type=float, default=10.0, help='weight for feature matching loss')
        parser.add_argument('--lambda_vgg', type=float, default=10.0, help='weight for vgg loss')

        return parser

if __name__ == '__main__':
    trainoptions = TrainOptions()
    opt = trainoptions.parse()